package com.pvdnc.encrypt;

import java.io.ByteArrayOutputStream;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Objects;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;

@SuppressWarnings("ALL")
public class RSA {
    public static final String RSA_ALGORITHM="RSA";

    private static Cipher getCipher(){
        try {
            return Cipher.getInstance(RSA_ALGORITHM);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
            e.printStackTrace();
            return null;
        }
    }

    public static byte[] encrypt(byte[] data, RSAPublicKey publicKey) throws IllegalStateException{
        //byte[] data=dataStr.getBytes();
        Cipher cipher=getCipher();
        Objects.requireNonNull(cipher);
        try {
            cipher.init(Cipher.ENCRYPT_MODE,publicKey);
        } catch (InvalidKeyException e) {
            throw new IllegalStateException(e);
        }
        int keySize=publicKey.getModulus().bitLength();
        byte[] finalData=splitCodec(cipher,Cipher.ENCRYPT_MODE,keySize,
                data);
        return finalData;
    }

    public static byte[] decrypt(byte[] data, RSAPrivateKey privateKey) throws IllegalStateException{
        //byte[] data= ObjectStream.toByte(encBase64);
        Cipher cipher=getCipher();
        Objects.requireNonNull(cipher);
        try {
            cipher.init(Cipher.DECRYPT_MODE,privateKey);
        } catch (InvalidKeyException e) {
            throw new IllegalStateException(e);
        }
        int keySize=privateKey.getModulus().bitLength();
        byte[] finalData=splitCodec(cipher,Cipher.DECRYPT_MODE,keySize,
                data);
        return finalData;
    }

    private static int getMaxBlock(int keySize,int opMode){
        int maxBlock=0;
        switch (opMode){//RSA是块加密
            case Cipher.DECRYPT_MODE:{
                maxBlock=keySize/8;
                break;
            }
            case Cipher.ENCRYPT_MODE:{
                maxBlock=keySize/8-11;
                break;
            }
            default:{
                throw new IllegalArgumentException("opMode must be "
                        +"Cipher.DECRYPT_MODE or Cipher.ENCRYPT_MODE");
            }
        }
        return maxBlock;
    }

    private static byte[] splitCodec(Cipher cipher,int opMode,int keySize,
                                     byte[] data) {
        int maxBlock = getMaxBlock(keySize, opMode);
        byte[] tempData;
        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
            int offset = 0, i = 0;
            while (data.length > offset) {
                if (data.length - offset > maxBlock)
                    tempData = cipher.doFinal(data, offset, maxBlock);
                else
                    tempData = cipher.doFinal(data, offset, data.length - offset);
                os.write(tempData, 0, tempData.length);
                i++;
                offset = i * maxBlock;
            }
            byte[] finalData = os.toByteArray();
            return finalData;
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
    }
}
